from unitree_legged_msgs.msg import MotorCmd
from unitree_legged_msgs.msg import MotorState
from unitree_legged_msgs.msg import LowState
import copy,rospy
class publisher:
    def __init__(self,joint_name_order,joint_command_dict) -> None:
        # rospy.init_node('joint_control', anonymous=True)
        self.publishers = {}
        self.motor_cmd_dict={}
        for joint_name in joint_name_order:        
            self.publishers[joint_name]= rospy.Publisher(joint_command_dict[joint_name], MotorCmd, queue_size = 1)
            motor_cmd = MotorCmd()
            # if "hip"  in joint_name:
            #     motor_cmd.mode = 0x0A                    
            #     motor_cmd.dq = 0.0
            #     motor_cmd.tau = 0.0
            #     motor_cmd.Kp = 60
            #     motor_cmd.Kd = 4
            #     self.motor_cmd_dict[joint_name]=motor_cmd

            # if "thigh"  in joint_name:
            #     motor_cmd.mode = 0x0A                    
            #     motor_cmd.dq = 0.0
            #     motor_cmd.tau = 0.0
            #     motor_cmd.Kp = 180
            #     motor_cmd.Kd = 8
            #     self.motor_cmd_dict[joint_name]=motor_cmd

            # if "calf"  in joint_name:
            #     motor_cmd.mode = 0x0A                    
            #     motor_cmd.dq = 0.0
            #     motor_cmd.tau = 0.0
            #     motor_cmd.Kp = 300
            #     motor_cmd.Kd = 15
            #     self.motor_cmd_dict[joint_name]=motor_cmd

            motor_cmd.mode = 0x0A                    
            motor_cmd.dq = 0.0
            motor_cmd.tau = 0.0
            motor_cmd.Kp = 85
            motor_cmd.Kd =2
        
            self.motor_cmd_dict[joint_name]=motor_cmd

    
    def publish_topics_smooth(self,joint_name_order,joint_angles_dict,duration):
        motor_cmd_dict_temp=copy.deepcopy(self.motor_cmd_dict)
        # smooth  的逻辑也需要改
        for i in range(1, int(duration) + 1):
            if not rospy.is_shutdown():
                percent = i / duration
                for joint_name in joint_name_order:
                    motor_cmd_dict_temp[joint_name].q = self.motor_cmd_dict[joint_name].q * (1 - percent) + joint_angles_dict[joint_name] * percent 
                    self.publishers[joint_name].publish(motor_cmd_dict_temp[joint_name])

        for joint_name in joint_name_order:
            self.motor_cmd_dict[joint_name].q=motor_cmd_dict_temp[joint_name].q

    def publish_topics(self,joint_name_order,joint_angles_dict):
        # smooth  的逻辑也需要改

        for joint_name in joint_name_order:
            self.motor_cmd_dict[joint_name].q =joint_angles_dict[joint_name]
            
            
        for joint_name in joint_name_order:
            self.publishers[joint_name].publish(self.motor_cmd_dict[joint_name])


    def set_position(self,joint_name_order,position_dict):
            for joint_name in joint_name_order:        
                self.motor_cmd_dict[joint_name].q=position_dict[joint_name]